import nibabel as nib
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import numpy as np
import os
from PIL import Image
from tqdm import tqdm
import cv2

def nii_to_png(nii_filename, output_folder):
    # 读取NIfTI文件
    nii_img = nib.load(nii_filename)
    nii_data = nii_img.get_fdata()

    # 创建输出文件夹
    os.makedirs(output_folder, exist_ok=True)

    # 将每个切片保存为PNG图像
    for i in tqdm(range(nii_data.shape[2]), ncols=50):
        # 获取当前切片
        slice_data = nii_data[:, :, i]

        # 创建PIL图像对象
        image = Image.fromarray(slice_data)

        # 将图像转换为灰度模式（如果需要）
        if image.mode != 'L':
            image = image.convert('L')

        # 保存为PNG文件
        png_filename = os.path.join(output_folder, f'slice_{i}.png')
        image.save(png_filename)


if __name__ == '__main__':
    nii_path = r'C:\Users\18049\Desktop\nii\3_mark_label\1-Pre\tumor-CHEN YI JIANG-M00043899-1-Pre.nii'
    output_folder = r"C:\develop\GraduationDesign\unet_seg\data\train\Ground"
    nii_to_png(nii_path, output_folder)
